import requests
import json
import re
from typing import List, Dict
import asyncio
import edge_tts
import pygame
import os
import threading
from pynput import keyboard

def clean_text(text: str) -> str:
    """清理文本，移除特殊字符和格式"""
    # 移除markdown语法
    text = re.sub(r'```[\s\S]*?```', '', text)  # 移除代码块
    text = re.sub(r'`[^`]*`', '', text)  # 移除行内代码
    text = re.sub(r'\*\*(.*?)\*\*', r'\1', text)  # 移除粗体
    text = re.sub(r'\*(.*?)\*', r'\1', text)  # 移除斜体
    text = re.sub(r'\[(.*?)\]\((.*?)\)', r'\1', text)  # 移除链接
    
    # 移除特殊符号和表情
    text = re.sub(r'[^\w\s\u4e00-\u9fff.,!?，。！？、:：;；]', '', text)
    
    # 格式化文本
    text = re.sub(r'\n\s*\n', '\n\n', text)  # 规范化段落间距
    text = re.sub(r'^\s+|\s+$', '', text, flags=re.MULTILINE)  # 移除行首尾空白
    
    return text.strip()

class DeepseekChat:
    def __init__(self, api_key: str):
        self.api_key = api_key
        self.base_url = "https://api.deepseek.com/v1/chat/completions"
        self.headers = {
            "Authorization": f"Bearer {api_key}",
            "Content-Type": "application/json"
        }
        self.conversation_history: List[Dict[str, str]] = []
        
        # 初始化TTS相关配置
        self.voices = {
            '1': 'zh-CN-XiaoxiaoNeural',    # 晓晓，女声
            '2': 'zh-CN-YunxiNeural',       # 云希，男声
            '3': 'zh-CN-YunjianNeural',     # 云健，男声
            '4': 'zh-CN-XiaoyiNeural',      # 晓伊，女声
            '5': 'zh-CN-YunyangNeural'      # 云扬，男声
        }
        self.current_voice = self.voices['1']
        self.is_playing = False
        self.enable_voice = True
        self.tts_task = None
        
        # 初始化键盘监听器
        self.listener = keyboard.Listener(on_press=self.on_press)
        self.listener.start()
    
    def on_press(self, key):
        """处理按键事件"""
        try:
            if hasattr(key, 'char'):
                if key.char in self.voices:  # 切换音色
                    self.current_voice = self.voices[key.char]
                    print(f"[已切换为{key.char}号音色]")
                elif key.char == 'v':  # 切换语音开关
                    self.enable_voice = not self.enable_voice
                    print(f"[语音已{'开启' if self.enable_voice else '关闭'}]")
                elif key.char == 's':  # 停止播放
                    if self.tts_task and not self.tts_task.done():
                        self.tts_task.cancel()
                        print("[已停止播放]")
        except Exception as e:
            pass
    
    async def _play_tts(self, text: str):
        """播放TTS语音"""
        try:
            communicate = edge_tts.Communicate(text, self.current_voice)
            self.is_playing = True
            
            async for chunk in communicate.stream():
                if chunk["type"] == "audio":
                    await asyncio.sleep(0)  # 让出控制权给事件循环
            
            self.is_playing = False
            
        except asyncio.CancelledError:
            self.is_playing = False
        except Exception as e:
            print(f"语音播放失败: {e}")
            self.is_playing = False
    
    def play_tts(self, text: str):
        """在新的事件循环中播放TTS语音"""
        if self.tts_task and not self.tts_task.done():
            self.tts_task.cancel()
        
        loop = asyncio.new_event_loop()
        self.tts_task = asyncio.run_coroutine_threadsafe(
            self._play_tts(text),
            loop
        )
    
    def chat(self, message: str) -> str:
        """发送消息并获取回复"""
        try:
            self.conversation_history.append({"role": "user", "content": message})
            
            data = {
                "model": "deepseek-chat",
                "messages": self.conversation_history,
                "stream": True,
                "temperature": 0.7,
                "max_tokens": 1000
            }
            
            response = requests.post(
                self.base_url,
                headers=self.headers,
                json=data,
                stream=True
            )
            
            if response.status_code != 200:
                return "抱歉，发生了错误，请稍后重试。"
            
            print("\nAI: ", end='', flush=True)
            full_response = ""
            
            for line in response.iter_lines():
                if line:
                    line = line.decode('utf-8')
                    if line.startswith('data: '):
                        try:
                            json_data = json.loads(line[6:])
                            if json_data.get('choices') and json_data['choices'][0].get('delta', {}).get('content'):
                                content = json_data['choices'][0]['delta']['content']
                                cleaned_content = clean_text(content)
                                if cleaned_content:
                                    print(cleaned_content, end='', flush=True)
                                    full_response += content
                        except json.JSONDecodeError:
                            continue
            
            print()  # 换行
            
            if full_response:
                self.conversation_history.append({"role": "assistant", "content": full_response})
                
                if self.enable_voice:
                    self.play_tts(clean_text(full_response))
            
            return full_response
                
        except Exception as e:
            print(f"发生错误: {e}")
            return "抱歉，发生了错误，请稍后重试。"
    
    def __del__(self):
        """清理资源"""
        try:
            if self.tts_task and not self.tts_task.done():
                self.tts_task.cancel()
        except:
            pass

def print_help():
    """打印帮助信息"""
    print("\n=== AI 助手使用说明 ===")
    print("\n基本操作：")
    print("- 直接输入文字进行对话")
    print("- 输入'help'显示帮助信息")
    print("- 输入'quit'退出程序")
    print("\n语音控制：")
    print("- 1-5: 切换不同音色")
    print("- v: 开启/关闭语音")
    print("- s: 停止播放")
    print("\n音色说明：")
    print("1: 晓晓（女声）")
    print("2: 云希（男声）")
    print("3: 云健（男声）")
    print("4: 晓伊（女声）")
    print("5: 云扬（男声）")
    print("\n当前设置：")
    print(f"- 语音：{'开启' if chat.enable_voice else '关闭'}")
    print(f"- 音色：{chat.current_voice}")
    print("="*30)

def main():
    global chat
    chat = DeepseekChat("sk-bab1e4f901344202b2f514efe1200128")
    print_help()
    
    while True:
        try:
            user_input = input("\n我: ").strip()
            
            if user_input.lower() == 'quit':
                print("再见！")
                break
            elif user_input.lower() == 'help':
                print_help()
                continue
            
            if not user_input:
                continue
            
            chat.chat(user_input)
            
        except KeyboardInterrupt:
            break
        except Exception as e:
            print(f"错误: {e}")
            continue

if __name__ == "__main__":
    main()
